#!/usr/bin/env python3
"""
Neural Trajectory 3D Time-Series Analyzer for Embodied World Models
分析稳定周期中的神经轨迹 - 3D时间序列可视化
"""

from grid_env_ideal_obs_repeat_task import *
from grid_agent import *
from checkpoint_utils import *
from maze_factory import *
from replay_config import *
import argparse
import json
import matplotlib.pyplot as plt
import time
from datetime import datetime
import numpy as np
from sklearn.decomposition import PCA
from mpl_toolkits.mplot3d import Axes3D
from mpl_toolkits.mplot3d.art3d import Poly3DCollection
from scipy.spatial import ConvexHull
import sys

@partial(jax.jit, static_argnums=(3,))
def model_forward(variables, state, x, model):
    """ forward pass of the model """
    return model.apply(variables, state, x)

@jit
def get_action(y):
    return jnp.argmax(y)
get_action_vmap = jax.vmap(get_action)

def load_task(pth = "./logs/task.json"):
    with open(pth, "r") as f:
        data = json.load(f)
        landscape = data["data"]
        state = data["state"]
        goal = data["goal"]
        print("Task loaded - Start:", state, "Goal:", goal)
    return landscape, state, goal

def progress_bar(current, total, width=50):
    """显示进度条"""
    percent = float(current) * 100 / total
    filled = int(width * current // total)
    bar = '█' * filled + '-' * (width - filled)
    sys.stdout.write(f'\r[{bar}] {percent:.1f}% ({current}/{total})')
    sys.stdout.flush()

def segment_neural_trajectories(neural_states, goal_reached_history):
    """根据目标到达历史分割神经轨迹为完整周期"""
    cycles = []
    cycle_start = 0
    
    for i, goal_reached in enumerate(goal_reached_history):
        if goal_reached:
            if i > cycle_start:
                cycle_neural_states = neural_states[cycle_start:i+1]
                cycles.append(cycle_neural_states)
                print(f"Cycle {len(cycles)}: steps {cycle_start}-{i} (length: {len(cycle_neural_states)})")
            cycle_start = i + 1
    
    return cycles

def analyze_neural_trajectories(cycles, num_cycles=10):
    """分析最后num_cycles个周期的神经轨迹"""
    if len(cycles) < num_cycles:
        print(f"Warning: Only {len(cycles)} cycles available, using all of them")
        selected_cycles = cycles
    else:
        selected_cycles = cycles[-num_cycles:]
        print(f"Analyzing last {num_cycles} cycles")
    
    # 将所有选定周期的神经状态合并
    all_neural_states = []
    cycle_labels = []
    
    for cycle_idx, cycle in enumerate(selected_cycles):
        all_neural_states.extend(cycle)
        cycle_labels.extend([cycle_idx] * len(cycle))
    
    all_neural_states = np.array(all_neural_states)
    cycle_labels = np.array(cycle_labels)
    
    print(f"Total neural states for analysis: {all_neural_states.shape}")
    
    # 执行PCA分析
    pca = PCA(n_components=3)
    pca_result = pca.fit_transform(all_neural_states)
    
    print(f"PCA explained variance ratio: {pca.explained_variance_ratio_}")
    print(f"Cumulative explained variance: {np.cumsum(pca.explained_variance_ratio_)}")
    
    return pca_result, cycle_labels, selected_cycles, pca

def visualize_neural_trajectories_3d_timeseries(pca_result, cycle_labels, selected_cycles, pca, all_cycles, perturbation_info=None):
    """3D时间序列可视化神经轨迹 - 横向时间轴上的垂直闭合极限环切片"""
    
    # 使用更大的figure和更高的DPI以获得最高分辨率的显示效果
    plt.rcParams['figure.dpi'] = 200  # 进一步增大显示DPI
    plt.rcParams['savefig.dpi'] = 800  # 进一步增大保存DPI
    fig = plt.figure(figsize=(24, 16))  # 增大figure尺寸配合高DPI
    ax = fig.add_subplot(111, projection='3d')
    
    # 绘制每个周期在3D时间序列中（横向排列）
    current_idx = 0
    
    # 动态计算时间轴间距，确保视觉上的明显分离
    num_cycles = len(selected_cycles)
    
    # 首先计算PCA结果的范围来确定合适的时间间距
    pca_range_y = np.max(pca_result[:, 0]) - np.min(pca_result[:, 0])  # PC1范围
    pca_range_z = np.max(pca_result[:, 1]) - np.min(pca_result[:, 1])  # PC2范围
    max_pca_range = max(pca_range_y, pca_range_z)
    
    # 时间间距设置为PCA范围的10倍，确保明显的视觉分离
    time_spacing = max_pca_range * 10.0
    print(f"Setting time_spacing to {time_spacing:.2f} (PCA range: {max_pca_range:.2f})")
    
    for cycle_idx, cycle in enumerate(selected_cycles):
        cycle_length = len(cycle)
        cycle_pca = pca_result[current_idx:current_idx + cycle_length]
        
        # 时间轴位置（现在用X轴表示时间）
        time_x = cycle_idx * time_spacing
        
        # 改进的颜色方案：使用viridis，避免过浅的颜色
        alpha = 1.0 - (cycle_idx / len(selected_cycles)) * 0.3  # 从1.0到0.7，保持可见性
        # 使用viridis颜色方案，颜色更鲜明，最后的颜色不会太浅
        color_ratio = 0.2 + (cycle_idx / len(selected_cycles)) * 0.7  # 从0.2到0.9，避免过暗和过浅
        color = plt.cm.viridis(color_ratio)
        
        # 创建闭合轨迹
        closed_cycle_pca = np.vstack([cycle_pca, cycle_pca[0:1]])
        
        # 添加时间轴坐标（X轴为时间，Y和Z轴为PC1和PC2）
        cycle_3d = np.column_stack([np.full(len(closed_cycle_pca), time_x),
                                  closed_cycle_pca[:, 0], closed_cycle_pca[:, 1]])
        
        # 绘制轨迹线条
        ax.plot(cycle_3d[:, 0], cycle_3d[:, 1], cycle_3d[:, 2], 
               color=color, linewidth=3, alpha=alpha)
        
        # 填充闭合区域（如果点数足够多）
        if cycle_length >= 3:
            try:
                # 计算2D凸包
                points_2d = cycle_pca[:, :2]
                hull = ConvexHull(points_2d)
                
                # 创建填充面片（垂直切片）
                hull_points_3d = []
                for simplex in hull.simplices:
                    triangle_3d = []
                    for point_idx in simplex:
                        triangle_3d.append([time_x, points_2d[point_idx, 0], points_2d[point_idx, 1]])
                    hull_points_3d.append(triangle_3d)
                
                # 添加半透明填充，颜色随时间变浅
                poly3d = Poly3DCollection(hull_points_3d, alpha=alpha*0.4, 
                                        facecolor=color, edgecolor='none')
                ax.add_collection3d(poly3d)
                
            except Exception as e:
                print(f"Warning: Could not create hull for cycle {cycle_idx}: {e}")
        
        # 标记起点和终点
        ax.scatter(time_x, cycle_pca[0, 0], cycle_pca[0, 1], 
                  color=color, s=120, marker='o', alpha=1.0, edgecolor='black')
        ax.scatter(time_x, cycle_pca[-1, 0], cycle_pca[-1, 1], 
                  color=color, s=120, marker='s', alpha=1.0, edgecolor='black')
        
        current_idx += cycle_length
    
    # ============ 扰动向量箭头已移除 ============
    # 移除扰动向量可视化功能
    
    # 手动设置坐标轴范围，强制时间轴占据主导地位
    time_range = (num_cycles - 1) * time_spacing
    ax.set_xlim([-time_spacing * 0.1, time_range + time_spacing * 0.1])  # 时间轴范围
    
    # 压缩Y和Z轴范围，让时间轴的间距更加突出
    y_margin = pca_range_y * 0.05
    z_margin = pca_range_z * 0.05
    ax.set_ylim([np.min(pca_result[:, 0]) - y_margin, np.max(pca_result[:, 0]) + y_margin])
    ax.set_zlim([np.min(pca_result[:, 1]) - z_margin, np.max(pca_result[:, 1]) + z_margin])
    
    # 强制设置坐标轴的纵横比，让时间轴拉伸
    ax.set_box_aspect([3, 1, 1])  # X轴（时间）是Y轴和Z轴的3倍长度
    
    # 设置简洁的坐标轴标签
    ax.set_xlabel('Episodes', fontsize=14, labelpad=20)  # 增加labelpad距离避免重叠
    ax.set_ylabel('PC-0', fontsize=14)
    ax.set_zlabel('PC-1', fontsize=14)
    ax.set_title('Neural Limit Cycles: 3D Time-Series View\n(Vertical slices across horizontal time axis)', fontsize=16)
    
    # 调整3D视角，让时间轴的间距更明显
    ax.view_init(elev=15, azim=20)  # 更低的仰角，更好地展示时间轴的延展
    
    # 设置坐标轴刻度
    time_ticks = [i * time_spacing for i in range(num_cycles)]
    ax.set_xticks(time_ticks)
    ax.set_xticklabels([f'{i+1}' for i in range(num_cycles)])  # 只显示数字，去掉"Cycle"文字
    
    # 去掉PC坐标轴上的数字刻度
    ax.set_yticks([])
    ax.set_zticks([])
    
    # 添加起点和终点的图例
    from matplotlib.lines import Line2D
    legend_elements = [
        Line2D([0], [0], marker='o', color='w', markerfacecolor='gray', markersize=10, 
               label='Episode Start', markeredgecolor='black'),
        Line2D([0], [0], marker='s', color='w', markerfacecolor='gray', markersize=10, 
               label='Episode End', markeredgecolor='black')
    ]
    ax.legend(handles=legend_elements, loc='upper right', fontsize=12)
    
    # 添加网格，但只在需要的地方
    ax.grid(True, alpha=0.3)
    
    # 显示分析信息文本
    fig.text(0.02, 0.98, f'Time spacing: {time_spacing:.2f}\nCycles: {num_cycles}', 
            verticalalignment='top', bbox=dict(boxstyle='round', facecolor='lightblue', alpha=0.8))
    
    # 移除tight_layout以避免布局压缩，使用subplots_adjust手动调整
    plt.subplots_adjust(left=0.05, right=0.95, top=0.95, bottom=0.05)
    
    # 保存文件时使用最高的DPI
    filename = './logs/neural_trajectories_3d_timeseries.png'
    plt.savefig(filename, dpi=800, bbox_inches='tight', facecolor='white')
    plt.show()

def main():
    """ Neural Trajectory 3D Time-Series Analysis Main Function """
    print("=" * 60)
    print("   Neural Trajectory 3D Time-Series Analyzer")
    print("   分析具身世界模型中的稳定周期神经轨迹 - 3D时间序列")
    print("=" * 60)
    
    # Parse arguments
    rpl_config = ReplayConfig()
    parser = argparse.ArgumentParser(description='Analyze neural trajectories in 3D time-series')
    parser.add_argument("--model_pth", type=str, default=rpl_config.model_pth)
    parser.add_argument("--task_pth", type=str, default=rpl_config.task_pth)
    parser.add_argument("--nn_size", type=int, default=rpl_config.nn_size)
    parser.add_argument("--nn_type", type=str, default=rpl_config.nn_type)
    parser.add_argument("--life_duration", type=int, default=1000)
    parser.add_argument("--num_cycles", type=int, default=10, help="Number of cycles to analyze")

    args = parser.parse_args()

    # Load model
    print("Loading model...")
    params = load_weights(args.model_pth)

    # Create agent
    if args.nn_type == "vanilla":
        model = RNN(hidden_dims=args.nn_size)
    elif args.nn_type == "gru":
        model = GRU(hidden_dims=args.nn_size)

    # Load task
    if os.path.isfile(args.task_pth):
        landscape, state, goal = load_task(pth=args.task_pth)
        landscape = [landscape]
    else:
        print("Error: Task file not found")
        return

    # Create environment
    print("Creating environment...")
    GE = GridEnv(landscapes=landscape, width=12, height=12, num_envs_per_landscape=1)
    GE.reset()
    
    # Set initial states
    GE.batched_states = GE.batched_states.at[0, 0].set(state[0])
    GE.batched_states = GE.batched_states.at[0, 1].set(state[1])
    GE.batched_goals = GE.batched_goals.at[0, 0].set(goal[0])
    GE.batched_goals = GE.batched_goals.at[0, 1].set(goal[1])
    GE.init_batched_states, GE.init_batched_goals = jnp.copy(GE.batched_states), jnp.copy(GE.batched_goals)
    GE.batched_goal_reached = batch_compute_goal_reached(GE.batched_states, GE.batched_goals)
    GE.last_batched_goal_reached = jnp.copy(GE.batched_goal_reached)
    GE.concat_obs = get_ideal_obs_vmap(GE.batched_envs, GE.batched_states, GE.batched_goals, GE.last_batched_goal_reached)

    concat_obs = GE.concat_obs
    rnn_state = model.initial_state(GE.num_envs)

    # Data collection with perturbation
    print("Starting neural state collection with perturbation...")
    neural_states_history = []
    goal_reached_history = []
    render_id = 0
    
    # 扰动相关变量
    perturbation_N = 100  # 最后N步中施加扰动
    perturbation_step = args.life_duration - np.random.randint(1, perturbation_N + 1)
    perturbation_vector = None
    perturbation_state_before = None
    perturbation_applied = False
    
    print(f"Perturbation will be applied at step {perturbation_step}")

    # Run simulation
    for step in range(args.life_duration):
        if step % 10 == 0:
            progress_bar(step, args.life_duration)

        # Model forward pass
        rnn_state, y1 = model_forward(params, rnn_state, concat_obs, model)
        
        batched_actions = get_action_vmap(y1)
        batched_goal_reached, concat_obs = GE.step(batched_actions, reset=True)

        # Record neural state and goal reached status
        neural_states_history.append(np.array(rnn_state[0]))
        goal_reached_history.append(bool(batched_goal_reached[render_id]))

        # 施加扰动（在记录神经状态之后，确保索引对齐）
        if step == perturbation_step and not perturbation_applied:
            perturbation_state_before = np.array(rnn_state[0])
            perturbation_vector = np.random.uniform(-0.3, 0.3, size=rnn_state[0].shape)
            rnn_state = rnn_state.at[0].add(perturbation_vector)
            perturbation_applied = True
            
            print(f"\nPerturbation applied at step {step}")
            print(f"Perturbation magnitude: {np.linalg.norm(perturbation_vector):.4f}")
            print(f"Perturbation will affect step {step + 1} neural state")

    print("\nData collection completed!")
    
    # Convert to numpy arrays
    neural_states_history = np.array(neural_states_history)
    goal_reached_history = np.array(goal_reached_history)
    
    print(f"Collected {len(neural_states_history)} neural states")
    print(f"Goal reached {np.sum(goal_reached_history)} times")

    # Segment into cycles
    print("\nSegmenting neural trajectories into cycles...")
    cycles = segment_neural_trajectories(neural_states_history, goal_reached_history)
    
    if len(cycles) == 0:
        print("Error: No complete cycles found!")
        return
    
    print(f"Found {len(cycles)} complete cycles")
    
    # Analyze trajectories
    print("\nAnalyzing neural trajectories...")
    pca_result, cycle_labels, selected_cycles, pca = analyze_neural_trajectories(cycles, args.num_cycles)
    
    # 准备扰动信息（保留用于数据分析，但不用于可视化）
    perturbation_info = None
    if perturbation_applied and perturbation_vector is not None:
        perturbation_info = (perturbation_step, perturbation_vector, perturbation_state_before)
    
    # Visualize in 3D time-series
    print("\nGenerating 3D time-series visualization...")
    visualize_neural_trajectories_3d_timeseries(pca_result, cycle_labels, selected_cycles, pca, cycles, perturbation_info)
    
    print("\nAnalysis completed! Results saved to ./logs/neural_trajectories_3d_timeseries.png")

if __name__ == "__main__":
    main()
